package com.xzx.flink.tableapi;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.functions.TableAggregateFunction;
import org.apache.flink.util.Collector;

import java.util.StringJoiner;

import static org.apache.flink.table.api.Expressions.$;
import static org.apache.flink.table.api.Expressions.call;

/**
 * 自定义UDF表聚合函数
 *
 * @version 1.0
 * @auther xinzhixuan
 * @date 2022/6/11 21:58
 */
public class TableAPI_08_UDF_TableAggregateFunction {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<User> stream = env.fromElements(
                new User("zhangsan", 20),
                new User("lisi", 30),
                new User("wangwu", 10));
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 目前 SQL 中没有直接使用表聚合函数的方式，所以需要使用 Table API 的方式来调用：
        tableEnv.createTemporaryView("my_table", stream, Schema.newBuilder()
                .column("name", DataTypes.STRING())
                .column("age", DataTypes.INT())
                .build());

        tableEnv.createTemporarySystemFunction("top2", Top2Function.class);

        Table resultTable = tableEnv.from("my_table")
                .flatAggregate(call("top2", $("name"), $("age")).as("user", "rank"))
                .select($("user"), $("rank"));

        tableEnv.toChangelogStream(resultTable).print();

        env.execute();

    }

    public static class User {
        public String name; // top1
        public Integer age; // top2

        public User() {
        }

        public User(String name, Integer age) {
            this.name = name;
            this.age = age;
        }

        @Override
        public String toString() {
            return new StringJoiner(", ", User.class.getSimpleName() + "[", "]")
                    .add("name='" + name + "'")
                    .add("age=" + age)
                    .toString();
        }
    }

    // 累加器类型定义
    public static class Top2Accumulator {
        public User first; // top1
        public User second; // top2
    }

    // 自定义聚合函数，输出为长整型的平均值,第一个泛型是最终输出的数据类型，这里输出一个二元组，<数据，及数据排名>
    public static class Top2Function extends TableAggregateFunction<Tuple2<String, Integer>, Top2Accumulator> {

        @Override
        public Top2Accumulator createAccumulator() {
            Top2Accumulator accumulator = new Top2Accumulator();
            accumulator.first = new User(null, Integer.MIN_VALUE);
            accumulator.second = new User(null, Integer.MIN_VALUE);
            return accumulator;
        }

        // 每来一个数据调用一次，判断是否更新累加器,第一个参数固定，后面的参数个数及类型不定
        public void accumulate(Top2Accumulator acc, String name, Integer age) {
            if (age > acc.first.age) {
                acc.second = acc.first;
                acc.first = new User(name, age);
            } else if (age > acc.second.age) {
                acc.second = new User(name, age);
            }
        }

        //所有输入行处理完成后，输出最终计算结果的方法。这个方法对应着 AggregateFunction中的 getValue()方法；
        // 输出(数值，排名)的二元组，输出两行数据,这个方法也是底层架构约定好的
        public void emitValue(Top2Accumulator acc, Collector<Tuple2<String, Integer>> out) {
            if (acc.first.age != Integer.MIN_VALUE) {
                out.collect(Tuple2.of(acc.first.toString(), 1));
            }
            if (acc.second.age != Integer.MIN_VALUE) {
                out.collect(Tuple2.of(acc.second.toString(), 2));
            }
        }
    }
}


